import numpy as np
import tqdm 
from tqdm import trange
from scipy.stats import multivariate_normal
import tensorflow as tf
import keras
#import numba as nb
import sys
import time
def attack(start_block,blocknum,trace_url,data_url,Z_url):
    model=[]
    for i in range(4):
        model.append(tf.keras.models.load_model(r"riscv-models\model_coeff{}.h5".format(i)))
        # 读取模型
    position=np.load(r"submit_code\start_idx_10.npy")
    # 切片起始位置
    result=[]
    guess=0
    right=0
    real=0
    count=0
    for n in trange(blocknum,desc="Iterating files"):
        trace_file=np.load(trace_url.format(n+start_block)).astype(np.float32)
        data_file=np.load(data_url.format(n+start_block))
        Z_file=np.load(Z_url.format(n+start_block))
        Z=Z_file["Z"]
        coefficient=data_file["allbytes"]
        index=data_file['index']

        coeff_index=[0,1,2,3]
        for i in trange(trace_file.shape[0]):

          
            for j in range(256):
                count+=1
                r=j//4
                if coefficient[i][j*8:j*8+8][:2]=="00":
                    real+=1
                start=position[r]
                end=start+1175

                piece=coefficient[i][j*8:j*8+8]
                piece_new=piece[6:8]+piece[4:6]+piece[2:4]+piece[0:2]
                t_slice=trace_file[i][start:end]
                Z_val=Z[i][index[i]][j]
                Z_LSB=(Z_val)&0xff
                if Z_LSB>128:
                    Z_LSB=-1*((~Z_LSB&0xff)+1)
                if j%4 in coeff_index and abs(Z_LSB)<=36:       ####
                    x = np.array(t_slice)
                    x = x.reshape(1, -1)
                    x=x.reshape([x.shape[0],x.shape[1],1])
                    # print(x.shape)

                    y=model[j%4].predict(x,verbose=0)
                    # if coefficient[i][j*8:j*8+8][:2]=="00":
                    #     print(y)
                    if y[0][0]>0.45:
                        guess+=1
                        result.append(np.array([n,i,j]))
                        if coefficient[i][j*8:j*8+8][:2]=="00":
                            right+=1
                #     elif y[0][0]<0.5 and coefficient[i][j*8:j*8+8][:2]!="00":
                #         right+=1
                # elif j%4 in coeff_index and abs(Z_LSB)>78 and coefficient[i][j*8:j*8+8][:2]!="00":
                #     right+=1
        print("Guess:{}, Right:{},  Real:{}".format(guess,right,real))
    print("Guess:{}, Right:{},  Real:{}".format(guess,right,real))
    return result




if __name__ == '__main__':
    func=1

    '''get result:func==1'''

    if func==1:
        tf.compat.v1.enable_eager_execution()
          
        for t in range(1):
            trace_url=r"D:\MLDSA_RISCV\traces_part{}.npy"
            data_url=r"Test_Metadata\metadata_part{}.npz"
            Z_url=r"Unpack\Test\Z_unpack_part{}.npz"
            from tensorflow.python.client import device_lib
            import os
            print(device_lib.list_local_devices()) # 查看本地可用设备
            # os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 忽略日志信息
            os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 指定使用哪个GPU，一个的情况下就不用了指定了
            tf.compat.v1.disable_eager_execution()  # 解决tensorflow版本2.0无法兼容版本1.0，不然sess.run()会出问题（RuntimeError: The Session graph is empty.  Add operations to the graph before calling run().）

            session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(log_device_placement=True))  # tensorflow2.0以后使用tf.compat.v1.Session(),之前直接使用tf.Session()

            blocknum=8
            start=0


            result=attack(start,blocknum,trace_url,data_url,Z_url)
            session.close()

            np.save(r"submit_code\result.npy",result)